import numpy as np
import matplotlib.pyplot as plt

fig,ax = plt.subplots()

data1 = np.loadtxt('SQETOTAL')
data2 = np.loadtxt('SQETOTAL_2')
data3 = np.loadtxt('SQETOTAL_3')

data = (data1+data2+data3)/3

sqe = data[:,2].reshape((4096,20),order='F')
sqe[sqe<=0] = 1e-20

np.savetxt('E.txt',data[0:4096,0])
np.savetxt('0p2_3.txt',sqe[:,3])
np.savetxt('0p4_3.txt',sqe[:,7])
np.savetxt('0p45_3.txt',sqe[:,8])
np.savetxt('0p5_3.txt',sqe[:,9])
np.savetxt('0p8_3.txt',sqe[:,15])

E = data[0:4096,0]


kB = 8.61733e-2 #meV/K
occ = 1./(np.exp(E/(kB*600))-1) + 1

chi = sqe/occ[:,None]
print(np.sum(chi))
np.savetxt('chi_ave',chi)
#-1,2
vmin=-2.5
vmax=-0.5
fig = ax.imshow(np.log10(chi),origin='lower',vmin=vmin,vmax=vmax,aspect='auto',cmap='viridis',interpolation='none')


#xpos = np.linspace(-1,19,2)
ax.set_xlim(0,19)
#xticks = np.linspace(0,1,2,dtype='int')
xpos = [-1,3,7,11,15,19]
xticks = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
ax.set_xticks(xpos)
ax.set_xticklabels(xticks)
ax.set_xlabel('0 0 L (r.l.u.)',fontsize=20)

#ypos = np.linspace(0,398,6)  # 398 20meV
#yticks = np.linspace(0,20,6,dtype='int')
#ax.set_ylim(0,398)
ypos = np.linspace(0,497,6)  # 497 25meV
yticks = np.linspace(0,25,6,dtype='int')
ax.set_ylim(0,497)
#ypos = np.linspace(0,596,7)
#yticks = np.linspace(0,30,7,dtype='int')
#ax.set_ylim(0,596)

ax.set_yticks(ypos)
ax.set_yticklabels(yticks)
ax.set_ylabel('Energy (meV)',fontsize=20)

ax.tick_params(which='both',labelsize=20)

#ax.axvline(9,0,1,color='white',ls='dashed',lw=2)
ax.text(3,490,'600K',ha='center',va='top',color='white',fontsize=40)

cbar = plt.colorbar(fig,ticks=np.arange(vmin,vmax+0.1,0.5))
cbar.ax.tick_params(labelsize=20)
cbar.ax.set_yticklabels(np.arange(vmin,vmax+0.1,0.5))
#plt.savefig('chi_ave.png',format='png',bbox_inches='tight')
#plt.savefig('chi_ave_05.pdf',format='pdf',bbox_inches='tight')
plt.show()
